# 在线文本行检测和识别
import json
import random
from common.detect import get_per_img_statistics_editdist, getpr
import cv2

from common.get_hw_pre import getresult
import numpy as np
import os

# 文字识别日志文件夹
diff_dir = "./indi"
os.makedirs(diff_dir, exist_ok=True)
# 标签格式是否是云平台,不然就是是线下的
label_style = True
if label_style:
    # 如果是云平台 需要指定图片文件夹
    img_dir = "E:/internship/dataset_test/cloud/indi_doc/images"
else:
    # 否则指定数据集文件夹
    data_dir = "E:/internship/dataset_test/8.25"

# 预测结果是 汉王格式 或者是百度格式
HW_format = False
if HW_format:
    # 在线还是离线, 在线需要内网，离线是使用保存的json预测结果
    on_line = False
    if not on_line:
        pre_json_dir = os.path.join("E:/internship/dataset_test/8.25","hw_bo_result.txt")
        pre_hw = dict()
        if not on_line:
            fpre_json_dir = open(pre_json_dir, "r", encoding="utf-8")
            pre_jsons = fpre_json_dir.readlines()
            for pre_json in pre_jsons:
                key = pre_json.split("\t")[0]
                key = key.replace("\\", "/")
                key = key.split("/")[-1]
                pre_hw[key] = pre_json.split("\t")[1]
else:
    fpaddle_sys_pre = os.path.join("E:/internship/dataset_test/cloud/", "indi_doc/system_results.txt")
    paddle_sys_pre = open(fpaddle_sys_pre, "r", encoding="utf-8")
    results = paddle_sys_pre.readlines()
    pre_paddle = dict()
    for result in results:
        label_temp = result.split("\t")[1]
        fn = result.split("\t")[0]
        pre_paddle[fn] = json.loads(label_temp)


# ===============================================================================
img_savedir = os.path.join(diff_dir,"str_det_result")
img_savedir = img_savedir.replace("\\","/")
os.makedirs(img_savedir,exist_ok=True)

if os.path.exists(os.path.join(diff_dir,"diff.txt")):
    os.remove(os.path.join(diff_dir,"diff.txt"))
if os.path.exists(os.path.join(diff_dir, "result-perimg.csv")):
    os.remove(os.path.join(diff_dir, "result-perimg.csv"))
fw = open(os.path.join(diff_dir, "result-perimg.csv"), 'a', encoding="utf-8")
fw.write("filename,det-percision,det-recall,rec-edit_score\n")
fw.close()
gdt_dict = dict()
# 文本行的真实数量
num_label = 0

# 不同的标签格式加载方式不一样
if not label_style:
    for line in open(os.path.join(data_dir, "Label.txt"), "r", encoding='utf-8'):
        strs = line.split("\t")
        filename = strs[0]
        # filename = filename.split(".")[0] + ".jpg"
        # path = os.path.join(data_dir, filename)
        path = os.path.join(data_dir, "../" + filename)
        path = path.replace("\\", "/")
        gdt_dict[path] = json.loads(strs[1])
else:

    label_dir = img_dir.replace("images","labels")
    fns = os.listdir(img_dir)
    for fn in fns:
        fgdt = open(os.path.join(label_dir, fn + ".txt"), "r", encoding="utf-8")
        fgdt_lines = fgdt.readlines()
        if len(fgdt_lines) > 1:
            print("!!!error : len(fgdt_lines)> 1")
        label_json = json.loads(fgdt_lines[0])
        if len(label_json) > 1:
            print("!!!error : len(label_json)> 1")
        path = os.path.join(img_dir, fn)
        path = path.replace("\\", "/")
        gdt_dict[path] = label_json[0]["Labels"]



sample_metrics = []
from tqdm import tqdm
for i, key in enumerate(tqdm(list(gdt_dict.keys()))):

    label = gdt_dict[key]
    # 文本行检测
    label_chars = []
    labe_boxes = []
    if not label_style :
        for j in range(len(label)):
            value = label[j]["transcription"]
            if ('picture'not in value) and ('para' not in value) and ('Para' not in value)and ('table' not in value) and (value.split("-")[0] != 'cell'):
                label_chars.append(value)
                xyxy = label[j]["points"]
                labe_boxes.append(xyxy)
    else:
        for current_label in label:
            value = current_label['Value']
            if ('picture'not in value) and ('para' not in value) and ('Para' not in value)and ('table' not in value) and (value.split("-")[0] != 'cell'):
                label_chars.append(value)
                xyxy = current_label['Points']
                if len(xyxy) == 4:
                    point1 = [round(xyxy[0]["X"]), round(xyxy[0]["Y"])]
                    point2 = [round(xyxy[1]["X"]), round(xyxy[1]["Y"])]
                    point3 = [round(xyxy[2]["X"]), round(xyxy[2]["Y"])]
                    point4 = [round(xyxy[3]["X"]), round(xyxy[3]["Y"])]
                    xyxy = [point1, point2, point3, point4]
                if len(xyxy) == 2:
                    point1 = [round(xyxy[0]["X"]), round(xyxy[0]["Y"])]
                    point2 = [round(xyxy[1]["X"]), round(xyxy[0]["Y"])]
                    point3 = [round(xyxy[1]["X"]), round(xyxy[1]["Y"])]
                    point4 = [round(xyxy[0]["X"]), round(xyxy[1]["Y"])]
                    xyxy = [point1, point2, point3, point4]
                labe_boxes.append([point1, point2, point3, point4])
    # pre:
    pre_chars = []
    pre_boxes = []
    if HW_format:
        if on_line:
            pre = getresult(key)
        else:
            hw_key = key.replace("\\","/")
            hw_key = hw_key.split("/")[-1]
            pre = pre_hw[hw_key]
            pre = json.loads(pre)
        paraphs = []
        if "Paraph" in pre.keys():
            paraphs = pre["Paraph"]
        all_pre_lines = []

        for para in paraphs:
            lines = para["line"]
            all_pre_lines = all_pre_lines + lines
            for line in lines:
                pre_chars.append(line["code"])
                coords = line["coords"]
                pre_boxes.append(coords)
    else :
        paddle_key = key.replace("\\", "/")
        paddle_key = paddle_key.split("/")[-1]
        results = pre_paddle[paddle_key]
        for result in results:
            value = result["transcription"]
            pre_chars.append(value)
            xyxy = result["points"]
            points = []
            # 4对坐标转为8个数值
            for xyxy_ in xyxy:
                for xyxy__ in xyxy_:
                    points.append(xyxy__)
            pre_boxes.append(points)
    fw = open(os.path.join(diff_dir,"diff.txt"), 'a', encoding="utf-8")
    fw.write(key + "\n")
    fw.close()

    current_metrics = get_per_img_statistics_editdist(pre_boxes, labe_boxes, 0.5, pre_chars, label_chars, diff_dir=diff_dir)
    true_positives_current, pred_scores_current, recong_score_current = [np.concatenate(x, 0) for x in list(zip(*current_metrics))]
    if len(true_positives_current) > 0:
        percision_current, recall_current = getpr(true_positives_current, len(labe_boxes))
    else:
        percision_current = -1
        recall_current = -1
    if percision_current > 0:
        sum_score = sum(recong_score_current)
        CAR = sum_score / len(recong_score_current)
        line_acc = np.sum(recong_score_current == 1) / len(recong_score_current)
    else:
        CAR = -1
        line_acc = -1
    fw = open(os.path.join(diff_dir, "result-perimg.csv"), 'a', encoding="utf-8")
    fw.write(key + "," + str(percision_current) + "," + str(recall_current) + "," + str(round(CAR, 4)) + "\n")
    fw.close()
    if (0 in current_metrics[0][0]) or len(labe_boxes) != len(pre_boxes):
        img = cv2.imread(key)
        if img is None:
            print("error cv2.imread() plase check img dir")
        imgcopy = img.copy()
        # 绘制预测框
        for prebox_index, pre_box in enumerate( pre_boxes):
            color = [0, 255, 0]
            aaa = current_metrics[0][0]
            tp_zeros_index = np.where(current_metrics[0][0] == 0)
            if len(tp_zeros_index[0]) > 0:
                if prebox_index in tp_zeros_index[0]:
                    color = [0, 0, 255]
            # pre_box 是8个数值
            for point_index in range(int(len(pre_box)/2)):
                offset = point_index * 2
                start = [pre_box[0 + offset],pre_box[1 + offset]]
                # 从最后一个点画到第一个点，数组内存溢出
                if 2 + offset > 7 :
                    offset = -2
                end = [pre_box[2 + offset],pre_box[3 + offset]]

                cv2.line(imgcopy, start, end, color, 2)
        for labe_box in labe_boxes:
            # labe_box 是4个二维数组
            for point_index in range(int(len(labe_box))):
                offset = point_index
                start = labe_box[0 + offset]
                if 1 + offset > 3 :
                    offset = -1
                end = labe_box[1 + offset]
                cv2.line(imgcopy, start, end, [255,0,0], 2)
        new_filename = os.path.join(img_savedir, key.split("/")[-1])
        saveresult = cv2.imwrite(new_filename, imgcopy)
        if not saveresult:
            print("!!!!error cv2.imwrite")

    sample_metrics += current_metrics

    num_label = num_label + len(labe_boxes)
true_positives, pred_scores, recong_score = [np.concatenate(x, 0) for x in list(zip(*sample_metrics))]

percision, recall = getpr(true_positives, num_label)
f_measure = (2 * percision * recall) / percision * recall

sum_score = sum(recong_score)
CAR = sum_score / len(recong_score)

line_acc = np.sum(recong_score == 1) / len(recong_score)

print("percision ", percision)
print("recall ", recall)
print("f_measure ", f_measure)
print("char acc ", CAR)
print("line acc ", line_acc)

fw = open(os.path.join(diff_dir, "result-perimg.csv"), 'a', encoding="utf-8")
fw.write("\n\ndet-percision,"+ str(percision) + "\n")
fw.write("det-recall,"+ str(recall)+ "\n")
fw.write("det-f_measure,"+ str(f_measure)+ "\n")
fw.write("rec-char acc,"+ str(CAR)+ "\n")
fw.write("rec-line acc,"+ str(line_acc)+ "\n")
fw.close()
